package com.wrapper.es;

import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.common.geo.GeoDistance;
import org.elasticsearch.common.unit.DistanceUnit;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.rescore.QueryRescoreMode;
import org.elasticsearch.search.rescore.QueryRescorerBuilder;
import org.elasticsearch.search.sort.*;
import org.springframework.data.domain.Sort;
import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.*;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.util.List;

import static org.springframework.util.CollectionUtils.isEmpty;

public class RequestFactory {

    // the default max result window size of Elasticsearch
    static final Integer INDEX_MAX_RESULT_WINDOW = 10_000;

    private final ElasticsearchConverter elasticsearchConverter;

    public RequestFactory(ElasticsearchConverter elasticsearchConverter) {
        this.elasticsearchConverter = elasticsearchConverter;
    }


    public SearchRequest searchRequest(Query query, @Nullable Class<?> clazz, IndexCoordinates index) {

        elasticsearchConverter.updateQuery(query, clazz);
        SearchRequest searchRequest = prepareSearchRequest(query, clazz, index);
        QueryBuilder elasticsearchQuery = getQuery(query);
        QueryBuilder elasticsearchFilter = getFilter(query);

        searchRequest.source().query(elasticsearchQuery);

        if (elasticsearchFilter != null) {
            searchRequest.source().postFilter(elasticsearchFilter);
        }

        return searchRequest;

    }

    private SearchRequest prepareSearchRequest(Query query, @Nullable Class<?> clazz, IndexCoordinates indexCoordinates) {

        String[] indexNames = indexCoordinates.getIndexNames();
        Assert.notNull(indexNames, "No index defined for Query");
        Assert.notEmpty(indexNames, "No index defined for Query");

        SearchRequest request = new SearchRequest(indexNames);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        sourceBuilder.version(true);
        sourceBuilder.trackScores(query.getTrackScores());
        if (hasSeqNoPrimaryTermProperty(clazz)) {
            sourceBuilder.seqNoAndPrimaryTerm(true);
        }

        if (query.getSourceFilter() != null) {
            SourceFilter sourceFilter = query.getSourceFilter();
            sourceBuilder.fetchSource(sourceFilter.getIncludes(), sourceFilter.getExcludes());
        }

        if (query.getPageable().isPaged()) {
            sourceBuilder.from((int) query.getPageable().getOffset());
            sourceBuilder.size(query.getPageable().getPageSize());
        } else {
            sourceBuilder.from(0);
            sourceBuilder.size(INDEX_MAX_RESULT_WINDOW);
        }

        if (!query.getFields().isEmpty()) {
            sourceBuilder.fetchSource(query.getFields().toArray(new String[0]), null);
        }

        if (query.getIndicesOptions() != null) {
            request.indicesOptions(query.getIndicesOptions());
        }

        if (query.isLimiting()) {
            // noinspection ConstantConditions
            sourceBuilder.size(query.getMaxResults());
        }

        if (query.getMinScore() > 0) {
            sourceBuilder.minScore(query.getMinScore());
        }

        if (query.getPreference() != null) {
            request.preference(query.getPreference());
        }

        request.searchType(query.getSearchType());

        prepareSort(query, sourceBuilder, getPersistentEntity(clazz));

        HighlightBuilder highlightBuilder = highlightBuilder(query);

        if (highlightBuilder != null) {
            sourceBuilder.highlighter(highlightBuilder);
        }

        if (query instanceof NativeSearchQuery) {
            prepareNativeSearch((NativeSearchQuery) query, sourceBuilder);
        }

        if (query.getTrackTotalHits() != null) {
            sourceBuilder.trackTotalHits(query.getTrackTotalHits());
        } else if (query.getTrackTotalHitsUpTo() != null) {
            sourceBuilder.trackTotalHitsUpTo(query.getTrackTotalHitsUpTo());
        }

        if (StringUtils.hasLength(query.getRoute())) {
            request.routing(query.getRoute());
        }

        TimeValue timeout = query.getTimeout();
        if (timeout != null) {
            sourceBuilder.timeout(timeout);
        }

        sourceBuilder.explain(query.getExplain());

        if (query.getSearchAfter() != null) {
            sourceBuilder.searchAfter(query.getSearchAfter().toArray());
        }

        query.getRescorerQueries().forEach(rescorer -> sourceBuilder.addRescorer(getQueryRescorerBuilder(rescorer)));

        request.source(sourceBuilder);
        return request;
    }

    private void prepareNativeSearch(NativeSearchQuery query, SearchSourceBuilder sourceBuilder) {

        if (!query.getScriptFields().isEmpty()) {
            for (ScriptField scriptedField : query.getScriptFields()) {
                sourceBuilder.scriptField(scriptedField.fieldName(), scriptedField.script());
            }
        }

        if (query.getCollapseBuilder() != null) {
            sourceBuilder.collapse(query.getCollapseBuilder());
        }

        if (!isEmpty(query.getIndicesBoost())) {
            for (IndexBoost indexBoost : query.getIndicesBoost()) {
                sourceBuilder.indexBoost(indexBoost.getIndexName(), indexBoost.getBoost());
            }
        }

        if (!isEmpty(query.getAggregations())) {
            for (AbstractAggregationBuilder<?> aggregationBuilder : query.getAggregations()) {
                sourceBuilder.aggregation(aggregationBuilder);
            }
        }

    }


    @Nullable
    private ElasticsearchPersistentEntity<?> getPersistentEntity(@Nullable Class<?> clazz) {
        return clazz != null ? elasticsearchConverter.getMappingContext().getPersistentEntity(clazz) : null;
    }

    private boolean hasSeqNoPrimaryTermProperty(@Nullable Class<?> entityClass) {

        if (entityClass == null) {
            return false;
        }

        if (!elasticsearchConverter.getMappingContext().hasPersistentEntityFor(entityClass)) {
            return false;
        }

        ElasticsearchPersistentEntity<?> entity = elasticsearchConverter.getMappingContext()
                .getRequiredPersistentEntity(entityClass);
        return entity.hasSeqNoPrimaryTermProperty();
    }

    private QueryRescorerBuilder getQueryRescorerBuilder(RescorerQuery rescorerQuery) {

        QueryBuilder queryBuilder = getQuery(rescorerQuery.getQuery());
        Assert.notNull("queryBuilder", "Could not build query for rescorerQuery");

        QueryRescorerBuilder builder = new QueryRescorerBuilder(queryBuilder);

        if (rescorerQuery.getScoreMode() != RescorerQuery.ScoreMode.Default) {
            builder.setScoreMode(QueryRescoreMode.valueOf(rescorerQuery.getScoreMode().name()));
        }

        if (rescorerQuery.getQueryWeight() != null) {
            builder.setQueryWeight(rescorerQuery.getQueryWeight());
        }

        if (rescorerQuery.getRescoreQueryWeight() != null) {
            builder.setRescoreQueryWeight(rescorerQuery.getRescoreQueryWeight());
        }

        if (rescorerQuery.getWindowSize() != null) {
            builder.windowSize(rescorerQuery.getWindowSize());
        }

        return builder;

    }

    @Nullable
    private QueryBuilder getQuery(Query query) {
        NativeSearchQuery searchQuery = (NativeSearchQuery) query;
        return searchQuery.getQuery();
    }

    @Nullable
    private QueryBuilder getFilter(Query query) {
        NativeSearchQuery searchQuery = (NativeSearchQuery) query;
        return searchQuery.getFilter();
    }

    @Nullable
    public HighlightBuilder highlightBuilder(Query query) {
        HighlightBuilder highlightBuilder = query.getHighlightQuery().map(HighlightQuery::getHighlightBuilder).orElse(null);

        if (highlightBuilder == null) {

            if (query instanceof NativeSearchQuery) {
                NativeSearchQuery searchQuery = (NativeSearchQuery) query;

                if (searchQuery.getHighlightFields() != null || searchQuery.getHighlightBuilder() != null) {
                    highlightBuilder = searchQuery.getHighlightBuilder();

                    if (highlightBuilder == null) {
                        highlightBuilder = new HighlightBuilder();
                    }

                    if (searchQuery.getHighlightFields() != null) {
                        for (HighlightBuilder.Field highlightField : searchQuery.getHighlightFields()) {
                            highlightBuilder.field(highlightField);
                        }
                    }
                }
            }
        }
        return highlightBuilder;
    }

    private void prepareSort(Query query, SearchSourceBuilder sourceBuilder,
                             @Nullable ElasticsearchPersistentEntity<?> entity) {

        if (query.getSort() != null) {
            query.getSort().forEach(order -> sourceBuilder.sort(getSortBuilder(order, entity)));
        }

        if (query instanceof NativeSearchQuery) {
            NativeSearchQuery nativeSearchQuery = (NativeSearchQuery) query;
            List<SortBuilder<?>> sorts = nativeSearchQuery.getElasticsearchSorts();
            if (sorts != null) {
                sorts.forEach(sourceBuilder::sort);
            }
        }
    }

    private SortBuilder<?> getSortBuilder(Sort.Order order, @Nullable ElasticsearchPersistentEntity<?> entity) {
        SortOrder sortOrder = order.getDirection().isDescending() ? SortOrder.DESC : SortOrder.ASC;

        if (ScoreSortBuilder.NAME.equals(order.getProperty())) {
            return SortBuilders //
                    .scoreSort() //
                    .order(sortOrder);
        } else {
            ElasticsearchPersistentProperty property = (entity != null) //
                    ? entity.getPersistentProperty(order.getProperty()) //
                    : null;
            String fieldName = property != null ? property.getFieldName() : order.getProperty();

            if (order instanceof GeoDistanceOrder) {
                GeoDistanceOrder geoDistanceOrder = (GeoDistanceOrder) order;

                GeoDistanceSortBuilder sort = SortBuilders.geoDistanceSort(fieldName, geoDistanceOrder.getGeoPoint().getLat(),
                        geoDistanceOrder.getGeoPoint().getLon());

                sort.geoDistance(GeoDistance.fromString(geoDistanceOrder.getDistanceType().name()));
                sort.ignoreUnmapped(geoDistanceOrder.getIgnoreUnmapped());
                sort.sortMode(SortMode.fromString(geoDistanceOrder.getMode().name()));
                sort.unit(DistanceUnit.fromString(geoDistanceOrder.getUnit()));
                return sort;
            } else {
                FieldSortBuilder sort = SortBuilders //
                        .fieldSort(fieldName) //
                        .order(sortOrder);

                if (order.getNullHandling() == Sort.NullHandling.NULLS_FIRST) {
                    sort.missing("_first");
                } else if (order.getNullHandling() == Sort.NullHandling.NULLS_LAST) {
                    sort.missing("_last");
                }
                return sort;
            }
        }
    }


}
